{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Replace StorageUnits with fundamental Links and Stores\n",
    "\n",
    "This notebook demonstrates how storage units can be replaced by more fundamental components, and how their parameters map to each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pypsa, os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from numpy.testing import assert_almost_equal, assert_array_almost_equal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define two functions we use in the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def replace_su(network, su_to_replace):\n",
    "    \"\"\"Replace the storage unit su_to_replace with a bus for the energy\n",
    "    carrier, two links for the conversion of the energy carrier to and from electricity,\n",
    "    a store to keep track of the depletion of the energy carrier and its\n",
    "    CO2 emissions, and a variable generator for the storage inflow.\n",
    "\n",
    "    Because the energy size and power size are linked in the storage unit by the max_hours,\n",
    "    extra functionality must be added to the LOPF to implement this constraint.\"\"\"\n",
    "\n",
    "    su = network.storage_units.loc[su_to_replace]\n",
    "\n",
    "    bus_name = \"{} {}\".format(su[\"bus\"], su[\"carrier\"])\n",
    "    link_1_name = \"{} converter {} to AC\".format(su_to_replace, su[\"carrier\"])\n",
    "    link_2_name = \"{} converter AC to {}\".format(su_to_replace, su[\"carrier\"])\n",
    "    store_name = \"{} store {}\".format(su_to_replace, su[\"carrier\"])\n",
    "    gen_name = \"{} inflow\".format(su_to_replace)\n",
    "\n",
    "    network.add(\"Bus\", bus_name, carrier=su[\"carrier\"])\n",
    "\n",
    "    # dispatch link\n",
    "    network.add(\n",
    "        \"Link\",\n",
    "        link_1_name,\n",
    "        bus0=bus_name,\n",
    "        bus1=su[\"bus\"],\n",
    "        capital_cost=su[\"capital_cost\"] * su[\"efficiency_dispatch\"],\n",
    "        p_nom=su[\"p_nom\"] / su[\"efficiency_dispatch\"],\n",
    "        p_nom_extendable=su[\"p_nom_extendable\"],\n",
    "        p_nom_max=su[\"p_nom_max\"] / su[\"efficiency_dispatch\"],\n",
    "        p_nom_min=su[\"p_nom_min\"] / su[\"efficiency_dispatch\"],\n",
    "        p_max_pu=su[\"p_max_pu\"],\n",
    "        marginal_cost=su[\"marginal_cost\"] * su[\"efficiency_dispatch\"],\n",
    "        efficiency=su[\"efficiency_dispatch\"],\n",
    "    )\n",
    "\n",
    "    # store link\n",
    "    network.add(\n",
    "        \"Link\",\n",
    "        link_2_name,\n",
    "        bus0=su[\"bus\"],\n",
    "        bus1=bus_name,\n",
    "        p_nom=su[\"p_nom\"],\n",
    "        p_nom_extendable=su[\"p_nom_extendable\"],\n",
    "        p_nom_max=su[\"p_nom_max\"],\n",
    "        p_nom_min=su[\"p_nom_min\"],\n",
    "        p_max_pu=-su[\"p_min_pu\"],\n",
    "        efficiency=su[\"efficiency_store\"],\n",
    "    )\n",
    "\n",
    "    if (\n",
    "        su_to_replace in network.storage_units_t.state_of_charge_set.columns\n",
    "        and (\n",
    "            ~pd.isnull(network.storage_units_t.state_of_charge_set[su_to_replace])\n",
    "        ).any()\n",
    "    ):\n",
    "        e_max_pu = pd.Series(data=1.0, index=network.snapshots)\n",
    "        e_min_pu = pd.Series(data=0.0, index=network.snapshots)\n",
    "        non_null = ~pd.isnull(\n",
    "            network.storage_units_t.state_of_charge_set[su_to_replace]\n",
    "        )\n",
    "        e_max_pu[non_null] = network.storage_units_t.state_of_charge_set[su_to_replace][\n",
    "            non_null\n",
    "        ]\n",
    "        e_min_pu[non_null] = network.storage_units_t.state_of_charge_set[su_to_replace][\n",
    "            non_null\n",
    "        ]\n",
    "    else:\n",
    "        e_max_pu = 1.0\n",
    "        e_min_pu = 0.0\n",
    "\n",
    "    network.add(\n",
    "        \"Store\",\n",
    "        store_name,\n",
    "        bus=bus_name,\n",
    "        e_nom=su[\"p_nom\"] * su[\"max_hours\"],\n",
    "        e_nom_min=su[\"p_nom_min\"] / su[\"efficiency_dispatch\"] * su[\"max_hours\"],\n",
    "        e_nom_max=su[\"p_nom_max\"] / su[\"efficiency_dispatch\"] * su[\"max_hours\"],\n",
    "        e_nom_extendable=su[\"p_nom_extendable\"],\n",
    "        e_max_pu=e_max_pu,\n",
    "        e_min_pu=e_min_pu,\n",
    "        standing_loss=su[\"standing_loss\"],\n",
    "        e_cyclic=su[\"cyclic_state_of_charge\"],\n",
    "        e_initial=su[\"state_of_charge_initial\"],\n",
    "    )\n",
    "\n",
    "    network.add(\"Carrier\", \"rain\", co2_emissions=0.0)\n",
    "\n",
    "    # inflow from a variable generator, which can be curtailed (i.e. spilled)\n",
    "    inflow_max = network.storage_units_t.inflow[su_to_replace].max()\n",
    "\n",
    "    if inflow_max == 0.0:\n",
    "        inflow_pu = 0.0\n",
    "    else:\n",
    "        inflow_pu = network.storage_units_t.inflow[su_to_replace] / inflow_max\n",
    "\n",
    "    network.add(\n",
    "        \"Generator\",\n",
    "        gen_name,\n",
    "        bus=bus_name,\n",
    "        carrier=\"rain\",\n",
    "        p_nom=inflow_max,\n",
    "        p_max_pu=inflow_pu,\n",
    "    )\n",
    "\n",
    "    if su[\"p_nom_extendable\"]:\n",
    "        ratio2 = su[\"max_hours\"]\n",
    "        ratio1 = ratio2 * su[\"efficiency_dispatch\"]\n",
    "\n",
    "        def extra_functionality(network, snapshots):\n",
    "            model = network.model\n",
    "            model.add_constraints(\n",
    "                model[\"Store-e_nom\"][store_name]\n",
    "                - model[\"Link-p_nom\"][link_1_name] * ratio1\n",
    "                == 0,\n",
    "                name=\"store_fix_1\",\n",
    "            )\n",
    "            model.add_constraints(\n",
    "                model[\"Store-e_nom\"][store_name]\n",
    "                - model[\"Link-p_nom\"][link_2_name] * ratio2\n",
    "                == 0,\n",
    "                name=\"store_fix_2\",\n",
    "            )\n",
    "\n",
    "    else:\n",
    "        extra_functionality = None\n",
    "\n",
    "    network.remove(\"StorageUnit\", su_to_replace)\n",
    "\n",
    "    return bus_name, link_1_name, link_2_name, store_name, gen_name, extra_functionality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, take an example from the git repo which has already been solved"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network_r = pypsa.examples.storage_hvdc(from_master=True)\n",
    "network_r.optimize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network = pypsa.examples.storage_hvdc(from_master=True)\n",
    "\n",
    "su_to_replace = \"Storage 0\"\n",
    "\n",
    "(\n",
    "    bus_name,\n",
    "    link_1_name,\n",
    "    link_2_name,\n",
    "    store_name,\n",
    "    gen_name,\n",
    "    extra_functionality,\n",
    ") = replace_su(network, su_to_replace)\n",
    "network.optimize(extra_functionality=extra_functionality)\n",
    "\n",
    "assert_almost_equal(network_r.objective, network.objective, decimal=2)\n",
    "assert_array_almost_equal(\n",
    "    network_r.storage_units_t.state_of_charge[su_to_replace],\n",
    "    network.stores_t.e[store_name],\n",
    ")\n",
    "assert_array_almost_equal(\n",
    "    network_r.storage_units_t.p[su_to_replace],\n",
    "    -network.links_t.p1[link_1_name] - network.links_t.p0[link_2_name],\n",
    ")\n",
    "\n",
    "# check optimised size\n",
    "assert_array_almost_equal(\n",
    "    network_r.storage_units.at[su_to_replace, \"p_nom_opt\"],\n",
    "    network.links.at[link_2_name, \"p_nom_opt\"],\n",
    ")\n",
    "assert_array_almost_equal(\n",
    "    network_r.storage_units.at[su_to_replace, \"p_nom_opt\"],\n",
    "    network.links.at[link_1_name, \"p_nom_opt\"]\n",
    "    * network_r.storage_units.at[su_to_replace, \"efficiency_dispatch\"],\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "",
   "language": "python",
   "name": ""
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "ed88634a96f0f44ddf87d723a7b512fbeabb17521926a161ee96c50fffea2b11"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
